package com.gitee.wsl.func.ext

import arrow.core.Either
import com.gitee.wsl.func.CoroutineCallSelfFun
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.selects.select
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract
import kotlin.coroutines.ContinuationInterceptor
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext


suspend infix fun <A,B> CoroutineCallSelfFun<A>.race(other : CoroutineCallSelfFun<B>)
    = raceN(fa = this, fb = other)


/**
 * Races the participants [fa], [fb] on the provided [CoroutineContext].
 * The winner of the race cancels the other participants.
 * Cancelling the operation cancels all participants.
 *
 * Coroutine context is inherited from a [CoroutineScope], additional context elements can be specified with [ctx] argument.
 * If the combined context does not have any dispatcher nor any other [ContinuationInterceptor], then [Dispatchers.Default] is used.
 * **WARNING** If the combined context has a single threaded [ContinuationInterceptor], this function will not run [fa] & [fb] in parallel.
 *
 * ```kotlin
 * import arrow.core.Either
 * import arrow.fx.coroutines.*
 * import kotlinx.coroutines.Dispatchers
 * import kotlinx.coroutines.awaitCancellation
 *
 * suspend fun main(): Unit {
 *   suspend fun loser(): Int =
 *     guaranteeCase({ awaitCancellation() }) { exitCase ->
 *       println("I can never win the race. Finished with $exitCase.")
 *     }
 *
 *   val winner = raceN(Dispatchers.IO, { loser() }, { 5 })
 *
 *   val res = when(winner) {
 *     is Either.Left -> "Never always loses race"
 *     is Either.Right -> "Race was won with ${winner.value}"
 *   }
 *   //sampleEnd
 *   println(res)
 * }
 * ```
 * <!--- KNIT example-race2-02.kt -->
 *
 * @param fa task to participate in the race
 * @param fb task to participate in the race
 * @return either [Either.Left] if [fa] won the race, or [Either.Right] if [fb] won the race.
 * @see raceN for a function that ensures it runs in parallel on the [Dispatchers.Default].
 */
@OptIn(ExperimentalContracts::class)
@Suppress("LEAKED_IN_PLACE_LAMBDA")
suspend inline fun <A, B> raceN(
    ctx: CoroutineContext = EmptyCoroutineContext,
    crossinline fa: suspend CoroutineScope.() -> A,
    crossinline fb: suspend CoroutineScope.() -> B
): Either<A, B> {
    contract {
        callsInPlace(fa, InvocationKind.AT_MOST_ONCE)
        callsInPlace(fb, InvocationKind.AT_MOST_ONCE)
    }
    return coroutineScope {
        val a = async(ctx) { fa() }
        val b = async(ctx) { fb() }
        select<Either<A, B>> {
            a.onAwait.invoke { Either.Left(it) }
            b.onAwait.invoke { Either.Right(it) }
        }.also {
            when (it) {
                is Either.Left -> b.cancelAndJoin()
                is Either.Right -> a.cancelAndJoin()
            }
        }
    }
}
